package com.eudon.ai.agent.etl.reader;

import com.eudon.ai.agent.common.MyIdGenerator;
import org.commonmark.node.*;
import org.commonmark.parser.Parser;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.Resource;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
/**
 * @author eudon
 * @description 由于Spring AI 框架并没有提供自定义ID generator的功能。这里直接重写MarkDocumentReader，然后自定义ID generator，后续框架更新了该类就废弃
 * @date 2025/8/20
 * @email syd19970616@gmail.com
 * @address <a href="https://gitee.com/eudon">Gitee</a>
 */
@Deprecated
public class MyMarkdownDocumentReader implements DocumentReader {
    private final Resource markdownResource;
    private final MarkdownDocumentReaderConfig config;
    private final Parser parser;

    public MyMarkdownDocumentReader(String markdownResource) {
        this((new DefaultResourceLoader()).getResource(markdownResource), MarkdownDocumentReaderConfig.defaultConfig());
    }

    public MyMarkdownDocumentReader(String markdownResource, MarkdownDocumentReaderConfig config) {
        this((new DefaultResourceLoader()).getResource(markdownResource), config);
    }

    public MyMarkdownDocumentReader(Resource markdownResource, MarkdownDocumentReaderConfig config) {
        this.markdownResource = markdownResource;
        this.config = config;
        this.parser = Parser.builder().build();
    }

    public List<Document> get() {
        try (InputStream input = this.markdownResource.getInputStream()) {
            Node node = this.parser.parseReader(new InputStreamReader(input));
            MyMarkdownDocumentReader.MyDocumentVisitor documentVisitor = new MyMarkdownDocumentReader.MyDocumentVisitor(this.config);
            node.accept(documentVisitor);
            return documentVisitor.getDocuments();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    static class MyDocumentVisitor extends AbstractVisitor {
        private final List<Document> documents = new ArrayList();
        private final List<String> currentParagraphs = new ArrayList();
        private final MarkdownDocumentReaderConfig config;
        private Document.Builder currentDocumentBuilder;

        MyDocumentVisitor(MarkdownDocumentReaderConfig config) {
            this.config = config;
        }

        @Override
        public void visit(org.commonmark.node.Document document) {
            this.currentDocumentBuilder = Document.builder().idGenerator(new MyIdGenerator());
            super.visit(document);
        }

        public void visit(Heading heading) {
            this.buildAndFlush();
            super.visit(heading);
        }

        public void visit(ThematicBreak thematicBreak) {
            if (this.config.horizontalRuleCreateDocument) {
                this.buildAndFlush();
            }

            super.visit(thematicBreak);
        }

        public void visit(SoftLineBreak softLineBreak) {
            this.translateLineBreakToSpace();
            super.visit(softLineBreak);
        }

        public void visit(HardLineBreak hardLineBreak) {
            this.translateLineBreakToSpace();
            super.visit(hardLineBreak);
        }

        public void visit(ListItem listItem) {
            this.translateLineBreakToSpace();
            super.visit(listItem);
        }

        public void visit(BlockQuote blockQuote) {
            if (!this.config.includeBlockquote) {
                this.buildAndFlush();
            }

            this.translateLineBreakToSpace();
            this.currentDocumentBuilder.metadata("category", "blockquote");
            super.visit(blockQuote);
        }

        public void visit(Code code) {
            this.currentParagraphs.add(code.getLiteral());
            this.currentDocumentBuilder.metadata("category", "code_inline");
            super.visit(code);
        }

        public void visit(FencedCodeBlock fencedCodeBlock) {
            if (!this.config.includeCodeBlock) {
                this.buildAndFlush();
            }

            this.translateLineBreakToSpace();
            this.currentParagraphs.add(fencedCodeBlock.getLiteral());
            this.currentDocumentBuilder.metadata("category", "code_block");
            this.currentDocumentBuilder.metadata("lang", fencedCodeBlock.getInfo());
            this.buildAndFlush();
            super.visit(fencedCodeBlock);
        }

        public void visit(Text text) {
            Node var3 = text.getParent();
            if (var3 instanceof Heading heading) {
                this.currentDocumentBuilder.metadata("category", "header_%d".formatted(heading.getLevel())).metadata("title", text.getLiteral());
            } else {
                this.currentParagraphs.add(text.getLiteral());
            }

            super.visit(text);
        }

        public List<Document> getDocuments() {
            this.buildAndFlush();
            return this.documents;
        }

        private void buildAndFlush() {
            if (!this.currentParagraphs.isEmpty()) {
                String content = String.join("", this.currentParagraphs);
                Document.Builder builder = this.currentDocumentBuilder.text(content);
                Map<String, Object> var10000 = this.config.additionalMetadata;
                Objects.requireNonNull(builder);
                var10000.forEach(builder::metadata);
                Document document = builder.build();
                this.documents.add(document);
                this.currentParagraphs.clear();
            }
            this.currentDocumentBuilder = Document.builder().idGenerator(new MyIdGenerator());
        }

        private void translateLineBreakToSpace() {
            if (!this.currentParagraphs.isEmpty()) {
                this.currentParagraphs.add(" ");
            }

        }
    }
}

